import cv2
import mediapipe as mp
import time
import matplotlib.pyplot as plt


def Mplot_landmarks(landmark_list, connections):
    if not landmark_list:
        return
    plt.figure(figsize=(10, 10))
    ax = plt.axes(projection="3d")
    ax.view_init(10, 10)
    plotted_landmarks = {}
    for idx, landmark in enumerate(landmark_list.landmark):
        if (landmark.HasField("visibility") and landmark.visibility < 0.5) or (
            landmark.HasField("presence") and landmark.presence < 0.5
        ):
            continue
        ax.scatter3D(xs=[-landmark.z], ys=[landmark.x], zs=[-landmark.y], c="r")
        plotted_landmarks[idx] = (-landmark.z, landmark.x, -landmark.y)
    if connections:
        num_landmarks = len(landmark_list.landmark)
        # Draws the connections if the start and end landmarks are both visible.
        for connection in connections:
            start_idx = connection[0]
            end_idx = connection[1]
            if not (0 <= start_idx < num_landmarks and 0 <= end_idx < num_landmarks):
                raise ValueError(
                    f"Landmark index is out of range. Invalid connection "
                    f"from landmark #{start_idx} to landmark #{end_idx}."
                )
            if start_idx in plotted_landmarks and end_idx in plotted_landmarks:
                landmark_pair = [
                    plotted_landmarks[start_idx],
                    plotted_landmarks[end_idx],
                ]
                ax.plot3D(
                    xs=[landmark_pair[0][0], landmark_pair[1][0]],
                    ys=[landmark_pair[0][1], landmark_pair[1][1]],
                    zs=[landmark_pair[0][2], landmark_pair[1][2]],
                    c="b",
                )
    plt.axis("equal")
    plt.show()


img = cv2.imread("pg.jpeg")
mpPos = mp.solutions.pose
pose = mpPos.Pose()
imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
width, height, _ = img.shape
results = pose.process(imgRGB)
if results.pose_world_landmarks:
    print("世界坐标：")  # 人体中心
    for id, lm in enumerate(results.pose_world_landmarks.landmark):
        print(id, lm.x, lm.y, lm.z)
if results.pose_landmarks:
    print("图像坐标：")  # 比例为x*width，y*height，z*depth，z为深度，为0时为图像外部
    for id, lm in enumerate(results.pose_landmarks.landmark):
        print(id, lm.x, lm.y, lm.z)
    mpDraw = mp.solutions.drawing_utils
    mpDraw.draw_landmarks(img, results.pose_landmarks, mpPos.POSE_CONNECTIONS)
    cv2.imshow("1", img)
    cv2.waitKey(0)


# 基于图像坐标轴  身高170  像素630  假设相机距离0.8m 630/170*80=296
var1 = results.pose_landmarks
for id, lm in enumerate(var1.landmark):
    lm.x = int(lm.x * height)
    lm.y = int(lm.y * width)
    lm.z = int(lm.z * 296)
    print(id, lm.x, lm.y, lm.z)
Mplot_landmarks(var1, mpPos.POSE_CONNECTIONS)

# 基于世界坐标轴   身高170cm  肚脐高度170*0.618≈105cm
var2 = results.pose_world_landmarks
for id, lm in enumerate(var2.landmark):
    lm.x = int(lm.x * 105)
    lm.y = int(lm.y * 105)
    lm.z = int(lm.z * 105)
    print(id, lm.x, lm.y, lm.z)
Mplot_landmarks(var2, mpPos.POSE_CONNECTIONS)


# mp.solutions.drawing_utils.plot_landmarks(
#     results.pose_world_landmarks, mpPos.POSE_CONNECTIONS
# )
